package com.wang.tool.trace;

import java.util.Map;
import java.util.Random;
import java.util.concurrent.atomic.AtomicInteger;

public class TraceContext {

    /**
     * callback来的调用不知道会被多少节点消费，所以给个极大的id
     * 100001->199999
     */
    public static final int THRESHOLD = 99999;
    private static final Random random = new Random();

    /**
     * not null
     */
    private String traceId;
    /**
     * not null
     */
    private String rpcId;

    private final AtomicInteger rpcSeq = new AtomicInteger(0);

    private Map<String, Object> attributes;

    public String buildReqId() {
        return traceId + "@" + rpcId;
    }


    public Object getAttribute(String key) {
        if (attributes == null) {
            return null;
        }
        return attributes.get(key);
    }

    public void putAttribute(String key, Object val) {
        if (attributes == null || key == null) {
            return;
        }
        if (val == null) {
            attributes.clear();
        } else {
            attributes.put(key, val);
        }
    }


    private TraceContext() {
    }

    /**
     * 根据入参创建上下文
     * @param traceId null -> return null
     * @param rpcId null -> 0
     * @param attributes null is valid
     * @return
     */
    public static TraceContext build(String traceId, String rpcId, Map<String, Object> attributes) {
        if (traceId == null) {
            return null;
        }
        TraceContext context = new TraceContext();
        context.traceId = traceId;
        context.rpcId = rpcId == null || rpcId.isEmpty() ? String.valueOf(context.rpcSeq.get()) : rpcId;
        context.attributes = attributes;
        return context;
    }

    /**
     * 根据入参创建回调场景的上下文
     * @param traceId null -> return null
     * @param rpcId null -> 0
     * @param attributes null is valid
     * @return
     */
    public static TraceContext buildCallbackTrace(String traceId, String rpcId, Map<String, Object> attributes) {
        if (traceId == null) {
            return null;
        }
        TraceContext context = new TraceContext();
        context.traceId = traceId;
        context.rpcId = rpcId == null || rpcId.isEmpty() ? String.valueOf(context.rpcSeq.get()) : (rpcId + "." + (random.nextInt(THRESHOLD) + THRESHOLD+2));
        context.attributes = attributes;
        return context;
    }

    public static TraceContext newTrace() {
        return newTrace(null);
    }

    public static TraceContext newTrace(Map<String, Object> attributes) {
        String traceId = TraceHolder.generateTraceId();
        if (traceId == null) {
            return null;
        }
        TraceContext context = new TraceContext();
        context.traceId = traceId;
        context.rpcId = String.valueOf(context.rpcSeq.get());
        context.attributes = attributes;
        return context;
    }

    public static TraceContext nextTrace(TraceContext context) {
        if (context == null) {
            String traceId = TraceHolder.generateTraceId();
            if (traceId == null) {
                return null;
            }
            context = new TraceContext();
            context.traceId = traceId;
            context.rpcId = String.valueOf(context.rpcSeq.get());
            return context;
        }
        return build(context.getTraceId(), context.rpcId + "." + context.rpcSeq.getAndIncrement(), context.getAttributes());
    }


    public String getTraceId() {
        return traceId;
    }

    public String getRpcId() {
        return rpcId;
    }

    public Map<String, Object> getAttributes() {
        return attributes;
    }
}
